#include "bs_cl.h"
#include "util.h"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>

//#define PTX_OUTPUT

void setup_cl(BS_test_t *t);
void setup_buffers(BS_test_t * t, cl_float(*gen)(cl_float low, cl_float high));
void warmup(BS_test_t * t, int num_iters);
bool run(BS_test_t * t, int num_iters);
bool check(BS_test_t *t);
	void BS_gold(BS_test_t *t);
		extern "C" void BlackScholesCPU(
		float *h_CallResult,
		float *h_PutResult,
		float *h_StockPrice,
		float *h_OptionStrike,
		float *h_OptionYears,
		float Riskfree,
		float Volatility,
		int optN
	);
	void print_arr(BS_test_t * t, cl_float *arr);
void post(BS_test_t *t, int run_num_iters, bool result);
void cleanup(BS_test_t *t);


/////////////////////////////////////////////////////

    const unsigned int   optionCount = 4000000;
    const float                    R = 0.02f;
    const float                    V = 0.30f;

/////////////////////////////////////////////////////

////////////////////////////////////////////////////////////////////////////////
// Helper function, returning uniformly distributed
// random float in [low, high] range
// Copyright 1993-2010 NVIDIA Corporation.  All rights reserved.
////////////////////////////////////////////////////////////////////////////////
cl_float RandFloat(cl_float low, cl_float high){
    cl_float t = (cl_float)rand() / (cl_float)RAND_MAX;
    return (1.0f - t) * low + t * high;
}

void setup_buffers(BS_test_t * t, cl_float(*gen)(cl_float low, cl_float high))
{
	// malloc host memory
	t->h_CallResult = (cl_float *) malloc (t->num * sizeof(cl_float));
	t->h_PutResult = (cl_float *) malloc (t->num * sizeof(cl_float));
	t->h_StockPrice = (cl_float *) malloc (t->num * sizeof(cl_float));
	t->h_OptionStrike = (cl_float *) malloc (t->num * sizeof(cl_float));
	t->h_OptionYears = (cl_float *) malloc (t->num * sizeof(cl_float));

	
	cl_int errcode_ret;
	// malloc device memory
	t->d_CallResult = clCreateBuffer(t->hContext, CL_MEM_WRITE_ONLY, t->num * sizeof(cl_float), NULL, &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateBuffer");
	t->d_PutResult = clCreateBuffer(t->hContext, CL_MEM_WRITE_ONLY, t->num * sizeof(cl_float), NULL, &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateBuffer");
	t->d_StockPrice = clCreateBuffer(t->hContext, CL_MEM_READ_ONLY, t->num * sizeof(cl_float), NULL, &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateBuffer");
	t->d_OptionStrike = clCreateBuffer(t->hContext, CL_MEM_READ_ONLY, t->num * sizeof(cl_float), NULL, &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateBuffer");
	t->d_OptionYears = clCreateBuffer(t->hContext, CL_MEM_READ_ONLY, t->num * sizeof(cl_float), NULL, &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateBuffer");

	// init host memory
	for (unsigned int i = 0; i < t->num; i++) {
		t->h_CallResult[i] = -1.0f;
		t->h_PutResult[i]  = -1.0f;
		t->h_StockPrice[i] = gen(5.0f, 30.0f);
		t->h_OptionStrike[i] = gen(1.0f, 100.0f);
		t->h_OptionYears[i] = gen(0.25f, 10.0f);
	}

	// copy to device memory
	printf("Copying input matrices to GPU...\n");
	errcode_ret = clEnqueueWriteBuffer(t->hCmdQueue, t->d_StockPrice, CL_TRUE, 0, t->num * sizeof(cl_float),
		t->h_StockPrice, 0, NULL, NULL);
	ERR_CHECK(errcode_ret, "clEnqueueWriteBuffer");
	errcode_ret = clEnqueueWriteBuffer(t->hCmdQueue, t->d_OptionStrike, CL_TRUE, 0, t->num * sizeof(cl_float),
		t->h_OptionStrike, 0, NULL, NULL);
	ERR_CHECK(errcode_ret, "clEnqueueWriteBuffer");
	errcode_ret = clEnqueueWriteBuffer(t->hCmdQueue, t->d_OptionYears, CL_TRUE, 0, t->num * sizeof(cl_float),
		t->h_OptionYears, 0, NULL, NULL);
	ERR_CHECK(errcode_ret, "clEnqueueWriteBuffer");
}

void setup_cl(BS_test_t *t)
{
	cl_int errcode_ret;

	// Get OpenCL platform count
	cl_uint NumPlatforms;
	clGetPlatformIDs (0, NULL, &NumPlatforms);

	// Get all OpenCL platform IDs
	cl_platform_id* PlatformIDs;
	PlatformIDs = new cl_platform_id[NumPlatforms];
	clGetPlatformIDs (NumPlatforms, PlatformIDs, NULL);

	// find NVIDIA & AMD platforms
	char cBuffer[1024];
	cl_int NvPlatform = -1;
	cl_int AMDPlatform = -1;
	for(cl_uint i = 0; i < NumPlatforms; ++i)
	{
		clGetPlatformInfo (PlatformIDs[i], CL_PLATFORM_NAME, 1024, cBuffer, NULL);
		printf("%s\n", cBuffer);
		if(strstr(cBuffer, "NVIDIA") != NULL) {
			NvPlatform = i;
		}
		else if (strstr(cBuffer, "AMD") != NULL) {
			AMDPlatform = i;
		}
	}

	// check for AMD and NVIDIA GPU devices
	cl_device_id cdDevice;
	cl_uint NvNumDevices = 0;
	cl_uint AMDNumDevices = 0;
	if (AMDPlatform != -1)
		clGetDeviceIDs(PlatformIDs[AMDPlatform], CL_DEVICE_TYPE_GPU, 0, NULL, &AMDNumDevices);
	if (NvPlatform != -1)
		clGetDeviceIDs(PlatformIDs[NvPlatform], CL_DEVICE_TYPE_GPU, 0, NULL, &NvNumDevices);

	// if there is an AMD GPU, take it, or take an NVIDIA GPU if it is there
	if (AMDNumDevices > 0)
		clGetDeviceIDs(PlatformIDs[AMDPlatform], CL_DEVICE_TYPE_GPU, 1, &cdDevice, NULL);
	else if (NvNumDevices > 0)
		clGetDeviceIDs(PlatformIDs[NvPlatform], CL_DEVICE_TYPE_GPU, 1, &cdDevice, NULL);
	else {
		fprintf(stderr, "could not find any GPU devices, exiting\n");
		delete [] PlatformIDs;
		exit(-1);
	}

	delete [] PlatformIDs;

	// get max work group size, just in case
	clGetDeviceInfo(cdDevice, CL_DEVICE_MAX_WORK_GROUP_SIZE, sizeof(size_t), &(t->maxBlockSize), NULL);
	
	//Create a context
	printf("Creating context for %s GPU...\n", (AMDNumDevices > 0) ? "AMD" : "NVIDIA");
	t->hContext = clCreateContext(NULL, 1, &cdDevice, NULL, NULL, &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateContext");

	size_t nContextDescriptorSize; 
	clGetContextInfo(t->hContext, CL_CONTEXT_DEVICES, 0, 0, &nContextDescriptorSize); 
	cl_device_id * aDevices = (cl_device_id *) malloc(nContextDescriptorSize); 
	clGetContextInfo(t->hContext, CL_CONTEXT_DEVICES, nContextDescriptorSize, aDevices, 0);

	// create a command queue for first
	// device the context reported  
	t->hCmdQueue = clCreateCommandQueue(t->hContext, aDevices[0], CL_QUEUE_PROFILING_ENABLE, &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateCommandQueue");
	free(aDevices);
	
	// create & compile Black-Scholes program
	cl_program hBSProg; 
	char * BSCode;
	size_t BSLen;
	
	printf("Compiling Black-Scholes program...\n");
	BSCode = oclLoadProgSource("bs_kernel.cl", "", &BSLen);
	hBSProg = clCreateProgramWithSource(t->hContext,1, (const char **)&BSCode, &BSLen, &errcode_ret); 
	ERR_CHECK(errcode_ret, "clCreateProgramWithSource BS");
	errcode_ret = clBuildProgram(hBSProg, 0, NULL, NULL, NULL, NULL);
	ERR_CHECK(errcode_ret, "clBuildProgram BS");
	free(BSCode);

#ifdef PTX_OUTPUT
	size_t progSize;
	errcode_ret = clGetProgramInfo(hBSProg, CL_PROGRAM_BINARY_SIZES, sizeof(size_t), &progSize, NULL);
	ERR_CHECK(errcode_ret, "clGetProgramInfo");
	char ** prog = (char **) malloc (sizeof(char **));
	prog[0] = (char *) malloc(progSize * sizeof(char *));
	errcode_ret = clGetProgramInfo(hBSProg, CL_PROGRAM_BINARIES, sizeof(char **) * progSize, prog, NULL);
	ERR_CHECK(errcode_ret, "clGetProgramInfo");
	FILE * f = fopen("bs_cl.ptx", "w");
	fprintf(f, "%s\n", prog[0]);
	fclose(f);
#endif

	// create BS kernel 
	printf("Creating Black-Scholes kernel...\n");
	t->hBSKernel = clCreateKernel(hBSProg, "BlackScholes", &errcode_ret);
	ERR_CHECK(errcode_ret, "clCreateKernel BlackScholes");
} 

cl_float gen0(cl_float low, cl_float high) {
	srand(5347);
	return RandFloat(low, high);
}

cl_float gen1(cl_float low, cl_float high) {
	return 5.0f;
}

void warmup(BS_test_t * t, int num_iters)
{
	t->run_num_iters = num_iters;
	launch_BS(t);
	printf("%d warmup iterations complete\n", num_iters);
}

bool run(BS_test_t * t, int num_iters)
{
	t->run_num_iters = num_iters;
	t->time = launch_BS(t);
	if (!check(t)) 
		return false;
	return true;
}

bool check(BS_test_t *t)
{
	float EPSILON = 2.0f;
	cl_int errcode_ret;

	BS_gold(t);
	cl_float * CallResult = (cl_float *) malloc (t->num * sizeof(cl_float));
	cl_float * PutResult = (cl_float *) malloc (t->num * sizeof(cl_float));
	errcode_ret = clEnqueueReadBuffer(t->hCmdQueue, t->d_CallResult, CL_TRUE, 0, t->num * sizeof(cl_float),
		CallResult, 0, NULL, NULL);
	ERR_CHECK(errcode_ret, "clEnqueueReadBuffer");
	errcode_ret = clEnqueueReadBuffer(t->hCmdQueue, t->d_PutResult, CL_TRUE, 0, t->num * sizeof(cl_float),
		PutResult, 0, NULL, NULL);
	ERR_CHECK(errcode_ret, "clEnqueueReadBuffer");
	//print_arr(t, CallResult);
	//print_arr(t, PutResult);

	for (unsigned int i = 0; i < t->num; i++) {
		cl_float hstCall = t->h_CallResult[i];
		cl_float hstPut = t->h_PutResult[i];
		cl_float devCall = CallResult[i];
		cl_float devPut = PutResult[i];
		if (abs(hstCall - devCall) >= EPSILON ||
			abs(hstPut - devPut) >= EPSILON ) {
				fprintf(stderr, "check failed at %d: [%f, %f] != [%f, %f]\n", 
					i,
					hstCall, hstPut,
					devCall, devPut);
				free(CallResult);
				free(PutResult);
				return false;
		}
	}
	free(CallResult);
	free(PutResult);
	return true;
}

void BS_gold(BS_test_t * t)
{
	BlackScholesCPU(
		t->h_CallResult,
		t->h_PutResult,
		t->h_StockPrice,
		t->h_OptionStrike,
		t->h_OptionYears,
		t->Riskfree,
		t->Volatility,
		t->num
	);
}

void print_arr(BS_test_t * t, cl_float *arr)
{
	for (unsigned int i = 0; i < t->num; i++) {
		if (i == 0)
			printf("[");
		printf("%3.0f", arr[i]);
		if (i < t->num-1)
			printf(", ");
		else
			printf("]\n");
	}
}

void post(BS_test_t * t, int run_num_iters, bool result)
{
	printf("----------------RESULTS----------------\n");
	printf("PASS? %s\n", result ? "YES" : "NO");
	printf("Size: %d work-items, work-group size of %d\n", t->gridSize, t->blockSize);
	printf("Time: %f\n", t->time / run_num_iters);
}

void output(BS_test_t *t, FILE *csv, int run_num_iters) {
	fprintf(csv, "%u, %u, %u, %1.9Lf\n",
		t->num,
		t->blockSize,
		t->gridSize,
		t->time / run_num_iters);
}

void cleanup(BS_test_t * t)
{
	cl_int errcode_ret;

	printf("Freeing Host and Device buffers...\n");
	free(t->h_CallResult);
	free(t->h_PutResult);
	free(t->h_StockPrice);
	free(t->h_OptionStrike);
	free(t->h_OptionYears);
	clFinish(t->hCmdQueue);

	errcode_ret = clReleaseMemObject(t->d_CallResult);
	ERR_CHECK(errcode_ret, "clReleaseMemObject d_CallResult");
	errcode_ret = clReleaseMemObject(t->d_PutResult);
	ERR_CHECK(errcode_ret, "clReleaseMemObject d_PutResult");
	errcode_ret = clReleaseMemObject(t->d_StockPrice);
	ERR_CHECK(errcode_ret, "clReleaseMemObject d_StockPrice");
	errcode_ret = clReleaseMemObject(t->d_OptionStrike);
	ERR_CHECK(errcode_ret, "clReleaseMemObject d_OptionStrike");
	errcode_ret = clReleaseMemObject(t->d_OptionYears);
	ERR_CHECK(errcode_ret, "clReleaseMemObject d_OptionYears");
	clFlush(t->hCmdQueue);
}

void test_BS (cl_uint sizeMin, cl_uint sizeMax, int warmup_num_iters, int run_num_iters, 
			  size_t blockMin, size_t blockMax, 
			  size_t gridMin, size_t gridMax,
			  char *filename, cl_float(*gen)(cl_float low, cl_float high))
{
	BS_test_t t;
	t.num = sizeMin;
	t.blockSize = blockMin;
	t.gridSize = gridMin;
	t.Riskfree = R;
	t.Volatility = V;
	bool result = false;

	FILE * csv = fopen(filename, "a");
	if (csv == NULL)
		fprintf(stderr, "can't open output file\n");

	if (csv != NULL)
		fprintf(csv, "size, block size, grid size, total time\n");

	setup_cl(&t);
	setup_buffers(&t, gen);

	if (t.maxBlockSize < blockMin)
		blockMin = t.maxBlockSize;
	if (t.maxBlockSize < blockMax)
		blockMax = t.maxBlockSize;

	printf("blockMIN: %d blockMAX: %d t.maxBlockSize: %d\n", blockMin, blockMax, t.maxBlockSize);

	warmup(&t, warmup_num_iters);

	for (t.num = sizeMin; t.num <= sizeMax; t.num <<= 1) {
		for (t.gridSize = gridMin; t.gridSize <= gridMax; t.gridSize += 16*1024) {
			for (t.blockSize = blockMin; t.blockSize <= blockMax; t.blockSize += 64) {
				if (t.gridSize % t.blockSize != 0)
					continue;
				setup_buffers(&t, gen);
				result = run(&t, run_num_iters);
				post(&t, run_num_iters, result);
				if (result && csv != NULL)
					output(&t, csv, run_num_iters);
				cleanup(&t);
			}
		}
	}

	if (csv != NULL)
		fclose(csv);
}

int main(int argc, char *argv[])
{
	char *filename = "output.csv";
	test_BS(optionCount, optionCount, 5, 50, 64, 512, 16*1024, 64*1024, filename, gen0);
	//test_BS(16, 16, 1, 5, 16, 16, 16, 16, filename, gen1);
	return 0;
}
